import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/v8/yolov8-omniFusion.yaml')
    # model.load('yolov8n.pt') # loading pretrain weights
    model.train(data='../dataset/data.yaml',
                cfg='ultralytics/cfg/cfgyolov8.yaml',
                cache=False,
                imgsz=640,
                epochs=100,
                batch=1,
                close_mosaic=10,
                workers=0,
                device='cuda',
                optimizer='auto', # using SGD
                # resume='', # last.pt path
                amp=False, # close amp
                # fraction=0.2,
                project='runs/train',
                name='yolov8base',
                )